import math

import dgl
import dgl.function as fn
import torch
import torch.nn as nn
import torch.nn.functional as F
from dgl.nn.pytorch import  GraphConv
from sklearn.metrics import precision_score
from torch.nn import init

import torch as th
from dgl import function as fn
from dgl._ffi.base import DGLError
from dgl.nn.pytorch import edge_softmax
from dgl.nn.pytorch.utils import Identity
from dgl.utils import expand_as_pair
from torch.nn import init


# 关系卷积层，考虑不同类型节点之间的关系特征。
class REConv(nn.Module):
    # 观察到在使用的时候，输入维度和输出维度都是节点类型数量num_type
    """
    Adapted from
    https://docs.dgl.ai/_modules/dgl/nn/pytorch/conv/gatconv.html#GATConv
    """
    # 初始化函数，定义输入参数。
    # in_feats：输入特征的维度。
    # out_feats：输出特征的维度。
    # norm：可选的归一化方法，支持 'none'、'both'、'right'、'left' 四种。
    # num_type：节点类型的数量，用于处理不同类型的节点。
    # weight：是否使用可学习的权重矩阵。
    # bias：是否使用偏置。
    # activation：激活函数。
    def __init__(self,
                 in_feats,
                 out_feats,
                 norm='both',
                 num_type=4,
                 weight=True,
                 bias=True,
                 activation=None):
        super(REConv, self).__init__()
        # 检查 norm 参数是否有效，如果不是支持的值，则抛出异常。
        if norm not in ('none', 'both', 'right', 'left'):
            raise DGLError('Invalid norm value. Must be either "none", "both", "right" or "left".'
                           ' But got "{}".'.format(norm))
        # 将输入参数 in_feats、out_feats 和 norm 存储为类的属性。
        self._in_feats = in_feats
        self._out_feats = out_feats
        self._norm = norm

        # 如果 weight 为 True，则定义一个可学习的权重矩阵 self.weight，其维度为 [in_feats, out_feats]。
        # 否则，不使用权重矩阵，并将其注册为 None。
        if weight:
            self.weight = nn.Parameter(th.Tensor(in_feats, out_feats))
        else:
            self.register_parameter('weight', None)

        # 如果 bias 为 True，则定义一个偏置向量 self.bias，其维度为 [out_feats]。
        # 否则，不使用偏置，并将其注册为 None。
        if bias:
            self.bias = nn.Parameter(th.Tensor(out_feats))
        else:
            self.register_parameter('bias', None)
        # 定义一个可学习的参数 self.weight_type，用来对每种节点类型的特征进行加权，初始值为全 1。
        # todo 这个地方注意一下有什么作用
        self.weight_type = nn.Parameter(th.ones(num_type))
        # 调用 reset_parameters 函数，初始化权重和偏置参数。
        self.reset_parameters()
        # 存储激活函数到类属性 _activation。
        self._activation = activation

    # 初始化权重和偏置参数。
    # 如果 self.weight 不为 None，使用 Xavier 初始化方法对权重进行初始化。
    # 如果 self.bias 不为 None，将偏置初始化为 0。
    def reset_parameters(self):
        if self.weight is not None:
            init.xavier_uniform_(self.weight)
        if self.bias is not None:
            init.zeros_(self.bias)

    # forward 是前向传播函数，接收三个参数：
    # graph：输入的图结构。
    # feat：输入节点的特征矩阵。
    # type_info：节点类型信息。
    # 使用 graph.local_scope()，确保对图结构的修改仅限于当前计算范围。
    # 过程：省略特征归一化，X=feat，'h'=X*W*weight_type，将源节点的'h'传递给目标节点的'h'，rst='h'，rst=激活函数(rst+b)
    def forward(self, graph, feat, type_info):
        with graph.local_scope():
            # 定义消息传递的聚合函数，将源节点的特征 'h' 复制为消息 'm'，用于消息传递。
            aggregate_fn = fn.copy_u('h', 'm')
            # 如果使用左归一化或双向归一化：
            # 计算出度（out_degrees），并限制最小值为 1，避免除以 0。
            # 如果 norm 是 'both'，则计算归一化因子 norm = deg^(-0.5)，否则使用 norm = 1.0 / deg。
            # 将归一化因子调整为与特征矩阵的维度相匹配，并对特征矩阵进行归一化。
            if self._norm in ['left', 'both']:
                # graph.out_degrees() 计算图中每个节点的出度，即每个节点的出边数量。
                # .float() 将出度数据转换为浮点数，以便后续的数学操作。
                # .clamp(min=1) 将所有出度的最小值限制为 1，避免出现除以 0 的情况（因为节点可能没有出边）。
                degs = graph.out_degrees().float().clamp(min=1)
                # 如果归一化模式为 'both'，则进行双边归一化，具体是对每个节点的出度开平方后取倒数，计算方式为 pow(degs, -0.5)。
                if self._norm == 'both':
                    norm = th.pow(degs, -0.5)
                else:
                    # 如果归一化模式为 'left'，则仅进行左归一化，计算方式为 1.0 / degs，即简单地对出度取倒数。
                    norm = 1.0 / degs
                # feat.dim() 返回特征张量的维度。
                # 这个操作将 norm 的形状扩展，确保它可以与 feat 进行广播相乘。假设 feat 是二维或多维张量，norm 需要扩展形状以匹配 feat 的维度。
                # 这个地方不太清楚是怎样计算的，不过无伤大雅
                shp = norm.shape + (1,) * (feat.dim() - 1)
                # 将 norm 重新调整为形状 shp，以便与 feat 兼容。
                norm = th.reshape(norm, shp)
                # 将特征矩阵 feat 按照归一化因子 norm 进行缩放，这种操作可以看作是对特征进行归一化，以考虑每个节点的连接强度。
                feat = feat * norm
            # 对输入的特征矩阵 feat 乘以权重矩阵 self.weight，完成特征的线性变换。
            feat = th.matmul(feat, self.weight)
            # 对节点特征进行加权，使用节点类型对应的 self.weight_type，并将结果保存到源节点的数据字段 'h'。
            graph.srcdata['h'] = feat * self.weight_type[type_info].reshape(-1, 1)
            # 使用消息传递机制：
            # 通过 aggregate_fn 将消息 'm'（即 'h'）从源节点传递到目标节点。
            # 使用 fn.sum 对消息进行求和，并将结果存储在目标节点的字段 'h' 中。
            graph.update_all(aggregate_fn, fn.sum(msg='m', out='h'))
            # 将目标节点的聚合结果存储在 rst 中。
            rst = graph.dstdata['h']

            # 如果使用右归一化或双向归一化：
            # 计算入度（in_degrees），限制最小值为 1。
            # 根据规范化的类型计算规范化因子 norm，并应用于目标节点的聚合结果 rst。
            if self._norm in ['right', 'both']:
                degs = graph.in_degrees().float().clamp(min=1)
                if self._norm == 'both':
                    norm = th.pow(degs, -0.5)
                else:
                    norm = 1.0 / degs
                shp = norm.shape + (1,) * (feat.dim() - 1)
                norm = th.reshape(norm, shp)
                rst = rst * norm
            # 如果使用偏置，则将偏置添加到聚合结果 rst 中。
            if self.bias is not None:
                rst = rst + self.bias
            # 如果定义了激活函数，则应用激活函数。
            if self._activation is not None:
                rst = self._activation(rst)

            # 返回计算结果 rst，即该层的输出。
            return rst


# AGTLayer 是一个自注意力（Self-Attention）机制的实现，它结合了关系特征（rh）来增强模型的表达能力。
class AGTLayer(nn.Module):
    # 初始化函数，定义输入参数：
    # embeddings_dimension：输入特征的维度。
    # nheads：多头注意力的头数。
    # att_dropout：注意力权重的 dropout 比例。
    # emb_dropout：输出特征的 dropout 比例。
    # temper：用于缩放注意力得分的温度系数。
    # rl：是否使用关系特征 rh。
    # rl_dim：关系特征的维度。
    # beta：权重系数，用于控制关系特征在注意力中的影响。
    def __init__(self, embeddings_dimension, nheads=2, att_dropout=0.5, emb_dropout=0.5, temper=1.0, rl=False, rl_dim=4, beta = 1):
        super(AGTLayer, self).__init__()

        # 将 nheads 和 embeddings_dimension 存储为类的属性。
        self.nheads = nheads
        self.embeddings_dimension = embeddings_dimension

        # 计算每个注意力头的维度，head_dim 是通过将总维度 embeddings_dimension 平均分配到每个注意力头数 nheads。
        self.head_dim = self.embeddings_dimension // self.nheads

        # 定义 LeakyReLU 作为激活函数，负斜率为 0.01，用于处理注意力得分。
        self.leaky = nn.LeakyReLU(0.01)
        # 将 temper、rl_dim 和 beta 存储为类属性，分别用于控制注意力的温度缩放、关系特征的维度、以及关系特征的权重。
        self.temper = temper
        self.rl_dim = rl_dim
        self.beta = beta

        # 定义两个线性层 linear_l 和 linear_r，分别用于将输入特征 h 映射到左、右查询特征（fl 和 fr），输出的维度是 head_dim * nheads。
        self.linear_l = nn.Linear(
            self.embeddings_dimension, self.head_dim * self.nheads, bias=False)
        self.linear_r = nn.Linear(
            self.embeddings_dimension, self.head_dim * self.nheads, bias=False)

        # 定义两个线性层 att_l 和 att_r，用于计算注意力得分的左、右部分，输出是标量，用于生成最终的注意力得分。
        self.att_l = nn.Linear(self.head_dim, 1, bias=False)
        self.att_r = nn.Linear(self.head_dim, 1, bias=False)

        # 如果 rl 为 True，则定义两个线性层 r_source 和 r_target，用于将关系特征 rh 映射到多头注意力的查询（r_q）和键（r_k），以便在注意力中引入关系特征。
        if rl:
            self.r_source = nn.Linear(rl_dim, rl_dim * self.nheads, bias=False)
            self.r_target = nn.Linear(rl_dim, rl_dim * self.nheads, bias=False)

        # 定义一个线性层 linear_final，用于将多头注意力的输出合并并映射回 embeddings_dimension 维度。
        self.linear_final = nn.Linear(
            self.head_dim * self.nheads, self.embeddings_dimension, bias=False)

        # 定义两个 Dropout 层，dropout1 用于注意力得分，dropout2 用于输出的特征表示。
        self.dropout1 = nn.Dropout(att_dropout)
        self.dropout2 = nn.Dropout(emb_dropout)

        # 定义一个 LayerNorm 层，用于对每个输入样本的特征进行标准化，保持数值稳定。
        self.LN = nn.LayerNorm(embeddings_dimension)

    # forward 是前向传播函数，接收两个输入参数：
    # h：节点的特征矩阵，维度为 [batch_size, seq_length, embeddings_dimension]。
    # rh：可选的关系特征矩阵，维度为 [batch_size, seq_length, rl_dim]，如果 rl 为 True，则需要提供。
    def forward(self, h, rh=None):
        # 获取输入特征 h 的 batch 大小，即 batch_size。
        batch_size = h.size()[0]
        # 将输入特征 h 分别通过 linear_l 和 linear_r 线性层，映射到查询特征 fl 和 fr，维度变化如下：
        # [batch_size, seq_length, embeddings_dimension] → [batch_size, seq_length, nheads, head_dim]
        # 然后使用 transpose(1, 2) 将头数维度提前，得到 [batch_size, nheads, seq_length, head_dim]。
        fl = self.linear_l(h).reshape(batch_size, -1, self.nheads, self.head_dim).transpose(1, 2)
        fr = self.linear_r(h).reshape(batch_size, -1, self.nheads, self.head_dim).transpose(1, 2)

        # 对 fl 和 fr 分别通过 att_l 和 att_r 线性层，计算左、右注意力得分，注意力得分表示为：
        # [batch_size, nheads, seq_length, 1] + [batch_size, nheads, 1, seq_length]
        # 最终得到 [batch_size, nheads, seq_length, seq_length] 的注意力得分矩阵。
        # self.att_r(self.leaky(fr)).permute(0, 1, 3, 2)：这一步对 self.att_r(self.leaky(fr)) 进行维度交换，使用 .permute(0, 1, 3, 2) 将张量的维度进行调整。permute 的作用是按照指定的顺序交换张量的维度：
        # 经过att_l和att_r后，维度变为 ([batch_size, nheads, seq_length, 1)。
        # 经过 .permute(0, 1, 3, 2) 后，维度变为 ([batch_size, nheads, 1, seq_length)，通常用于在后续的矩阵运算中进行广播相加或者内积操作。
        # 最终，经过 att_l 和 att_r 的结果相加，得到 score。这个 score 可能代表注意力得分，用于后续的注意力机制中。
        # score 的维度是 [batch_size, nheads, seq_length, seq_length]
        score = self.att_l(self.leaky(fl)) + self.att_r(self.leaky(fr)).permute(0, 1, 3, 2)

        # 如果 rh 不是 None，则引入关系特征 rh：
        # 将 rh 分别通过 r_source 和 r_target 映射到关系特征的键 r_k 和查询 r_q。
        # 计算关系特征的注意力得分 score_r，其维度为 [batch_size, nheads, seq_length, seq_length]。
        # 将 score_r 乘以系数 beta 后加到原始的 score 中，增强了关系特征在注意力中的作用。
        if rh is not None:
            # self.r_source 是一个线性层（例如 nn.Linear），对输入特征 rh 进行线性变换，生成键向量 r_k。
            # reshape(batch_size, -1, self.nheads, self.rl_dim) 将输出重塑为 [batch_size, seq_length, nheads, rl_dim] 的形状。
            # transpose(1, 2) 交换第二和第三维度，得到 [batch_size, nheads, seq_length, rl_dim]，将多头注意力的头部放在维度 1。
            r_k = self.r_source(rh).reshape(batch_size, -1, self.nheads, self.rl_dim).transpose(1,2)
            # self.r_target 是另一个线性层，对同样的 rh 进行线性变换，生成查询向量 r_q。
            # reshape 后同样得到 [batch_size, seq_length, nheads, rl_dim] 的形状。
            # permute(0, 2, 3, 1) 交换维度，得到 [batch_size, nheads, rl_dim, seq_length]，准备进行矩阵乘法。
            r_q = self.r_target(rh).reshape(batch_size, -1, self.nheads, self.rl_dim).permute(0, 2, 3, 1)
            # r_k @ r_q 是张量的矩阵乘法，r_k 的形状是 [batch_size, nheads, seq_length, rl_dim]，r_q 的形状是 [batch_size, nheads, rl_dim, seq_length]。
            # 矩阵乘法后得到 score_r 的形状为 [batch_size, nheads, seq_length, seq_length]，即计算了不同位置间的相似性或相关性。
            score_r = r_k @ r_q
            # 将先前的 score 加上 self.beta * score_r，这里 self.beta 是一个标量系数，用于调整 score_r 的影响。最终的 score 是更新后的注意力得分。
            score = score + self.beta * score_r

        # 使用温度系数 temper 对注意力得分进行缩放，调节注意力机制的敏感度。
        score = score / self.temper

        # 对注意力得分应用 softmax 函数，使其在最后一个维度上归一化，从而得到注意力权重。[batch_size, nheads, seq_length, seq_length]
        score = F.softmax(score, dim=-1)
        # 对注意力权重进行 dropout，随机丢弃一些权重，防止过拟合。
        score = self.dropout1(score)

        # 计算上下文表示 context，即使用注意力权重 score 对 fr 进行加权求和，维度为 [batch_size, nheads, seq_length, head_dim]。
        context = score @ fr

        # 将上下文表示 context 进行变换，将头数维度 nheads 合并回原特征维度，得到形状为 [batch_size, seq_length, embeddings_dimension] 的特征。
        h_sa = context.transpose(1,2).reshape(batch_size, -1, self.head_dim * self.nheads)
        # 将上下文表示通过 linear_final 线性层，映射回原始维度，并进行 dropout。
        fh = self.linear_final(h_sa)
        fh = self.dropout2(fh)

        # 使用 LayerNorm 对残差连接的结果进行归一化。h + fh 是残差结构，即原始输入特征 h 加上自注意力输出 fh。
        h = self.LN(h + fh)

        # 返回经过注意力机制后的特征 h，其维度为 [batch_size, seq_length, embeddings_dimension]。
        return h


# HINormer 是一个基于图神经网络 (GCN)、关系卷积 (REConv) 和自注意力机制 (AGTLayer) 的混合模型，用于处理异构图的节点分类任务。模型通过多层的图卷积、关系卷积和自注意力机制对节点的特征进行提取和聚合，最终输出节点的分类结果。
class HINormer(nn.Module):
    # g: 图数据，通常是 DGLGraph，表示图结构。
    # num_class: 分类任务的类别数量。
    # input_dimensions: 每种节点类型的输入特征维度列表。
    # embeddings_dimension: 模型中每一层的输出特征维度。
    # num_layers: 自注意力层的数量。
    # num_gnns: 图卷积层的数量。
    # nheads: 自注意力机制中多头注意力的头数。
    # dropout: dropout 比例，用于防止过拟合。
    # temper: 自注意力机制的温度参数。
    # num_type: 节点类型的数量。
    # beta: 自注意力层中关系特征的权重参数。
    def __init__(self, g, num_class, input_dimensions, embeddings_dimension=64, num_layers=8, num_gnns=2, nheads=2, dropout=0,  temper=1.0, num_type=4, beta=1):

        super(HINormer, self).__init__()

        # 初始化类的成员变量，用于存储图数据、嵌入维度、层数、类别数等信息。
        self.g = g
        self.embeddings_dimension = embeddings_dimension
        self.num_layers = num_layers
        self.num_class = num_class
        self.num_gnns = num_gnns
        self.nheads = nheads

        # 使用 ModuleList 创建一个包含多个线性层的列表，每个线性层将不同类型节点的输入特征（input_dimensions）映射到 embeddings_dimension 维度。这用于不同类型节点的特征变换。
        self.fc_list = nn.ModuleList([nn.Linear(in_dim, embeddings_dimension) for in_dim in input_dimensions])
        # 存储 dropout 参数，用于控制模型的正则化。
        self.dropout = dropout
        # 分别创建 GCNLayers、RELayers 和 GTLayers，用于存储多层的图卷积层 (GraphConv)、关系卷积层 (REConv) 和自注意力层 (AGTLayer)。
        self.GCNLayers = torch.nn.ModuleList()
        self.RELayers = torch.nn.ModuleList()
        self.GTLayers = torch.nn.ModuleList()
        # 循环 num_gnns 次，分别在 GCNLayers 和 RELayers 中添加 GraphConv 和 REConv 层：
        # GraphConv：标准的图卷积层，使用 ReLU 作为激活函数，将输入特征映射到 embeddings_dimension 维度。
        # REConv：关系卷积层，考虑不同类型节点之间的关系特征。
        for layer in range(self.num_gnns):
            self.GCNLayers.append(GraphConv(
                self.embeddings_dimension, self.embeddings_dimension, activation=F.relu))

            self.RELayers.append(REConv(num_type, num_type, activation=F.relu, num_type=num_type))

        # 循环 num_layers 次，向 GTLayers 添加自定义的自注意力层 AGTLayer：
        # AGTLayer 通过多头注意力机制，处理节点之间的交互，并通过参数 rl=True 允许使用关系特征
        for layer in range(self.num_layers):
            self.GTLayers.append(
                AGTLayer(self.embeddings_dimension, self.nheads, self.dropout, self.dropout, temper=temper, rl=True, rl_dim=num_type, beta=beta))

        # 定义一个 Dropout 层，防止模型过拟合。
        self.Drop = nn.Dropout(self.dropout)

        # 定义一个线性层 Prediction，用于将节点的最终特征映射到类别空间，输出每个类别的概率分布。
        self.Prediction = nn.Linear(embeddings_dimension, num_class)

    # features_list: 输入的节点特征列表，不同类型的节点有不同的特征。
    # seqs: 序列信息，用于选择特定节点的特征进行处理。
    # type_emb: 节点类型的嵌入矩阵，一个单位矩阵（one-hot编码）
    # node_type: 每个节点对应的类型信息。
    # norm: 可选参数，是否对最终输出进行归一化。
    def forward(self, features_list, seqs, type_emb, node_type, norm=False):
        h = []
        # 对每种类型的节点特征 features_list 使用线性层 fc_list 进行变换，生成相同维度的特征，将结果存储在 h 中。
        for fc, feature in zip(self.fc_list, features_list):
            # 先统一h中列(特征)的维度，即h中的每个矩阵的维度为 [节点数，embeddings_dimension]
            h.append(fc(feature))
        # 将不同类型节点的特征沿着维度 0 进行拼接，得到所有节点的特征表示 gh。[总节点数，embeddings_dimension]
        gh = torch.cat(h, 0)
        # 从 type_emb 中选择与当前节点类型 node_type 对应的嵌入，得到节点的类型嵌入 r。
        # node_type 是一个张量，形状为 [num_nodes]，表示每个节点的类型。它的值是每个节点类型的索引（整数），例如：node_type[i] = t 表示第 i 个节点的类型是 t，其中 t 是从 0 到 num_node_types - 1 的整数。
        r = type_emb[node_type]
        # 依次对图卷积层和关系卷积层进行前向传播：
        # GCNLayers[layer](self.g, gh)：在图 g 上对节点特征 gh 进行图卷积操作，更新节点表示。
        # Drop(gh)：对更新后的节点表示进行 dropout 操作。
        # RELayers[layer](self.g, r, node_type)：使用 REConv 在图上更新类型嵌入 r，同时考虑节点的类型信息 node_type。
        for layer in range(self.num_gnns):
            gh = self.GCNLayers[layer](self.g, gh)
            gh = self.Drop(gh)
            r = self.RELayers[layer](self.g, r, node_type)
        # 根据序列信息 seqs，从 gh 和 r 中选择出目标节点的特征表示和类型嵌入。
        # gh 的维度为 [num_nodes, embeddings_dimension]，表示图中所有节点的特征矩阵。
        # seqs 的维度为 [batch_size, seq_length]，表示一组节点序列的索引。
        # gh[seqs] 会提取这些索引位置的节点特征，返回的张量 h 的维度是 [batch_size, seq_length, embeddings_dimension]，
        # 其中：batch_size 是批次中的序列数量。seq_length 是每条序列中的节点数量。embeddings_dimension 是每个节点的特征维度。
        h = gh[seqs]
        # h 的维度是 [batch_size, seq_length, num_type]
        r = r[seqs]
        # 依次通过自注意力层 GTLayers 进行特征的聚合与更新：
        # h：节点的特征表示。
        # rh：节点的关系特征，作为额外输入增强特征的表达能力。
        for layer in range(self.num_layers):
            h = self.GTLayers[layer](h, rh = r)
        # 对经过多层自注意力更新后的节点特征 h 通过线性层 Prediction 映射到类别空间，得到分类的输出 output。
        # h[:, 0, :]之后，形状将变为 [batch_size, embeddings_dimension]
        output = self.Prediction(h[:, 0, :])
        # 如果 norm 为 True，则对输出结果进行归一化，防止数值过大影响模型稳定性。
        if norm:
            output = output / (torch.norm(output, dim=1, keepdim=True)+1e-12)
        # 返回模型的最终输出，即每个节点对应的类别预测结果。
        return output
